import json
import sys
from typing import List, Optional, Dict, Any, Tuple, Union
from uuid import UUID
import Ex_Tools
import Prompt
import PromtpTpl
import Tools
from langchain.memory import ConversationTokenBufferMemory
from langchain.tools.render import render_text_description
from langchain_core.callbacks import BaseCallbackHandler
from langchain_core.language_models import BaseChatModel
from langchain_core.output_parsers import PydanticOutputParser, StrOutputParser
from langchain_core.outputs import GenerationChunk, ChatGenerationChunk, LLMResult
from langchain_core.prompts import PromptTemplate
from langchain_core.tools import StructuredTool


from pydantic import BaseModel, Field, ValidationError


class AIAgent:
    def __init__(self,
                 llm: BaseChatModel = ChatOpenAI(
                     model="gpt-4.0-turbo",
                     temperature=0.3,
                     model_kwargs={
                         "seed": 42
                     }
                 ),
                 tools=None,
                 prompt: str = "",
                 final_prompt: str = "",
                 max_thought_step: Optional[int] = 10,
                 ):
        if tools is None:
            tools = []
        self.llm = llm
        self.tools = tools
        self.final_prompt = PromptTemplate.from_template(final_prompt)
        self.max_thought_steps = max_thought_step  # 最多思考步骤
        self.prompt = self.__init__prompt(prompt)
        self.llm_chain = self.prompt | self.llm | StrOutputParser()
        self.verbose_printer = Ex_Tools.MyPrintHandler()
        self.output_parser = PydanticOutputParser(
            pydantic_object=Ex_Tools.Action
        )

    @staticmethod
    def __chinese_friendly(string) -> str:
        lines = string.split("\n")
        for i, line in enumerate(lines):
            if line.startwith('{') and line.endswith('}'):
                try:
                    lines[i] = json.dumps(json.loads(line), ensure_ascii=False)
                except:
                    print("警报警报！出现错误啦！！！！")
        return "\n".join(lines)

    def run(self, task_description):
        thought_step_count = 0
        agent_memory = ConversationTokenBufferMemory(
            llm=self.llm,
            max_token_limit=4000
        )
        agent_memory.save_context(
            {"input": "\ninit"},
            {"output": "\n开始"}
        )
        while thought_step_count < self.max_thought_steps:
            print(f">>>>Round: {thought_step_count}<<<<")
            action, response = self.__step(
                task_description=task_description,
                memory=agent_memory
            )

            # 如果是结束指令，执行最后一步
            if action.name == "FINISH":
                break

            # 执行动作
            observation = self.__exec_action(action)
            print(f"----\nObservation:\n{observation}")

            # 更新记忆
            self.__update_memory(agent_memory, response, observation)

            thought_step_count += 1

        if thought_step_count >= self.max_thought_steps:
            # 如果思考步数达到上限，返回错误信息
            reply = "抱歉，我没能完成您的任务。"
        else:
            # 否则，执行最后一步
            final_chain = self.final_prompt | self.llm | StrOutputParser()
            reply = final_chain.invoke({
                "task_description": task_description,
                "memory": agent_memory
            })

        return reply

    def __step(self, task_description, memory) -> Tuple[Ex_Tools.Action, str]:

        """执行一步思考"""
        response = ""
        for s in self.llm_chain.stream({
            "task_description": task_description,
            "memory": memory
        }, config={
            "callbacks": [
                self.verbose_printer
            ]
        }):
            response += s

        action = self.output_parser.parse(response)
        return action, response

    def __exec_action(self, action: Ex_Tools.Action) -> str:
        observation = "没有找到工具"
        for tool in self.tools:
            if tool.name == action.name:
                try:
                    # 执行工具
                    observation = tool.run(action.args)
                except ValidationError as e:
                    # 工具的入参异常
                    observation = (
                        f"Validation Error in args: {str(e)}, args: {action.args}"
                    )
                except Exception as e:
                    # 工具执行异常
                    observation = f"Error: {str(e)}, {type(e).__name__}, args: {action.args}"

        return observation

    @staticmethod
    def __update_memory(agent_memory, response, observation):
        agent_memory.save_context(
            {"input": response},
            {"output": "\n返回结果:\n" + str(observation)}
        )

    def __init__prompt(self, prompt):
        return PromptTemplate.from_template(prompt).partial(
            tools=render_text_description(self.tools),
            format_instructions=self.__chinese_friendly(
                self.output_parser.get_format_instructions()
            )
        )


if __name__ == "__main__":
    my_agent = AIAgent(
        tools=Tools.tools,
        prompt=Prompt.prompt_text,
        final_prompt=Prompt.final_prompt,
    )

    task = "帮我买24年6月1日早上去上海的火车票"
    reply = my_agent.run(task)
    print(reply)
